# Date: 2025/03/25
import time
import numpy as np
from coppeliasim_zmqremoteapi_client import RemoteAPIClient
import os
import math
import basedriver
import roboticstoolbox as rtb
from motion_planning import Motion_Planning
from spatialmath import *
from armkinematics import Jaka,get_trajectory,get_object_pose

# 初始化手指关节角度（以度为单位）
initial_finger_joint_angles_deg = [0, 0, -20, -20]
finger_joint_angles_rad = [math.radians(angle) for angle in initial_finger_joint_angles_deg]
left_wheel_speed, right_wheel_speed = basedriver.inverse(0, 0)

robot_arm = Jaka()

# 定义手指和机械臂关节名称列表
finger_joint_names = ['finger01_joint01', 'finger01_joint02', 'finger01_joint03', 'finger01_joint04']
arm_joint_names = ['Joint01', 'Joint02', 'Joint03', 'Joint04', 'Joint05', 'Joint06']

# 所有关节名称列表
all_joint_names = [
    *[f'finger0{i}_joint0{j}' for i in range(1, 5) for j in range(1, 5)],
    'LM_motor', 'RM_motor', 'Joint06', 'Joint05', 'Joint04', 'Joint03', 'Joint02', 'Joint01'
]

# 所有关节句柄字典
joint_handles = {}

def initialize_simulation():
    """初始化仿真环境"""
    client = RemoteAPIClient()
    sim = client.require('sim')
    sim.stopSimulation()
    print('成功与仿真软件建立通信！', flush=True)
    scene_path = 'd:/VrepModels/v02(whole body)(6).ttt'
    sim.loadScene(scene_path)
    print('成功加载场景！', flush=True)
    for joint_name in all_joint_names:
        handle = sim.getObject('/'+joint_name)
        joint_handles[joint_name] = handle
    sim.setStepping(True)
    print('开始仿真！', flush=True)
    sim.startSimulation()
    sim.setJointTargetVelocity(joint_handles['LM_motor'], 0)
    sim.setJointTargetVelocity(joint_handles['RM_motor'], 0)
    return sim

def get_initial_angles(sim, joint_names):
    """获取关节初始角度"""
    initial_angles_deg = []
    for joint_name in joint_names:
        current_angle = sim.getJointPosition(joint_handles[joint_name])
        current_angle_deg = math.degrees(current_angle)
        initial_angles_deg.append(current_angle_deg)

    return initial_angles_deg

def plan_arm_trajectory(start_angles_rad , robot_model):
    """规划机械臂轨迹"""
    via_points = np.array([
        start_angles_rad,
        [-2.75575965, -2.91763972, -0.87049957,  2.21734296, -1.57079633,
       -1.18496332]#新生成的测试轨迹
    ])

    precatch_point = np.array([-2.75575965, -2.91763972, -0.87049957,  2.21734296, -1.57079633,
       -1.18496332])

    print("start angles")
    print(start_angles_rad,flush = True)

    arm_trajectory = rtb.jtraj(start_angles_rad,precatch_point, 50)

    return arm_trajectory

def execute_arm_trajectory(sim, arm_trajectory):
    """执行机械臂轨迹"""
    
    for point_index, point in enumerate(arm_trajectory.q):
        for joint_index, joint_name in enumerate(arm_joint_names):
            sim.setJointTargetPosition(joint_handles[joint_name], point[joint_index])
            # sim.step()

        while True:
            # 检查每个关节是否到达目标位置
            if all(np.isclose(sim.getJointPosition(joint_handles[joint_name]), point[joint_index], atol=1e-3) 
                   for joint_index, joint_name in enumerate(arm_joint_names)):
                break
            sim.step()
    # sim.step()

def plan_and_execute_finger_trajectory(sim, start_angles_rad, end_angles_rad, num_points):
    """规划并执行手指轨迹"""
    motion_planner = Motion_Planning()
    finger_trajectory = motion_planner.jtraj_cal(start_angles_rad, end_angles_rad, num_points)
    finger_trajectory = finger_trajectory.astype(np.float64)
    for trajectory_point in range(num_points):
        for finger_index in range(1, 5):
            for joint_index in range(4):
                joint_name = f'finger0{finger_index}_joint0{joint_index+1}'
                sim.setJointTargetPosition(joint_handles[joint_name], finger_trajectory[trajectory_point][joint_index])
        sim.setJointTargetVelocity(joint_handles['Joint06'], 0)
        sim.setJointTargetVelocity(joint_handles['LM_motor'], 0)
        sim.setJointTargetVelocity(joint_handles['RM_motor'], 0)
        sim.step()
        # print(f"执行手指轨迹点 {trajectory_point+1}/{num_points}", flush=True)


def print_format_transformation_matrix(sim,object_name):

    object_handle = sim.getObject(object_name)
    matrix = sim.getObjectMatrix(object_handle)
    # 初始化一个4x4的单位矩阵
    transformation_matrix = [
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1]
    ]
    
    # 将输入的12元素矩阵填充到4x4矩阵中
    for i in range(3):
        for j in range(4):
            if j == 3:
                transformation_matrix[i][j] = matrix[i * 4 + j] * 1000
            else:
                transformation_matrix[i][j] = matrix[i * 4 + j]

    for row in transformation_matrix:
        print(' '.join(
            [f"{elem:.3e}" if len(str(elem).replace('.', '').replace('e', '').replace('+', '').replace('-', '')) > 4 else str(elem) for elem in row]
        ), flush=True)
    
    return transformation_matrix


def format_transformation_matrix(matrix):
    # 初始化一个4x4的单位矩阵
    transformation_matrix = [
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1]
    ]
    
    # 将输入的12元素矩阵填充到4x4矩阵中
    for i in range(3):
        for j in range(4):
            transformation_matrix[i][j] = matrix[i * 4 + j]
    
    return transformation_matrix


def main():
    sim = initialize_simulation()
    try:
        # 获取手指和机械臂初始角度
        current_finger_angles_deg = get_initial_angles(sim, finger_joint_names)
        current_arm_angles_deg = get_initial_angles(sim, arm_joint_names)
        
        # 转换为弧度
        start_finger_angles_rad = np.deg2rad(current_finger_angles_deg)
        start_arm_angles_rad = np.deg2rad(current_arm_angles_deg)
        end_finger_angles_rad = np.array(finger_joint_angles_rad)

        print("初始时刻基坐标系的位姿矩阵：", flush=True)
        hand_baselink_matrix = print_format_transformation_matrix(sim,'/Link_06_visual')
        
        # 规划并执行机械臂轨迹
        # arm_trajectory = plan_arm_trajectory(start_arm_angles_rad, robot_arm)
        arm_trajectory1,arm_trajectory2 ,arm_trajectory3= get_trajectory(get_object_pose())
        execute_arm_trajectory(sim, arm_trajectory1)
        print("到达预抓取姿态")

        execute_arm_trajectory(sim, arm_trajectory2)
        print("到达抓取姿态")

        print("末尾时刻基坐标系的位姿矩阵：", flush=True)
        hand_baselink_matrix = print_format_transformation_matrix(sim,'/Link_06_visual')

        print("物体位姿矩阵：", flush=True)
        object_matrix = print_format_transformation_matrix(sim,'/Cuboid')

        print("手臂运动完毕！",flush=True)

        plan_and_execute_finger_trajectory(sim, start_finger_angles_rad, end_finger_angles_rad, 19)
        print("抓取完成！", flush=True)

        execute_arm_trajectory(sim, arm_trajectory3)
        print("到达放置姿态")    
        


        # 持续仿真直到用户中断
        while True:
            sim.step()
    except KeyboardInterrupt:
        print('程序被用户终止！')
    finally:
        sim.stopSimulation()
        print('仿真已停止。')

if __name__ == "__main__":
    main()